import traceback
import argparse
import datetime
import json
import cv2
import os

__CLASS__ = ['__background__', 'kpts']   # class dictionary, background must be in first index.

# phase_folder  images annotations 这三个文件夹需要对应,用于生成train.json&val.json，每次生成一个
def argparser():
    parser = argparse.ArgumentParser("define argument parser for pycococreator!")
    # 根目录
    parser.add_argument("-r", "--root_path", default="data/images_crop", help="path of root directory")

    # 两个文件夹：train, val
    parser.add_argument("-p", "--phase_folder", default=['val'], help="datasets path of [train, val, test]")

    # 判断是否有关键点
    parser.add_argument("-po", "--have_points", default=True, help="if have points we will deal with it!")

    # 图片文件夹
    parser.add_argument("-im", "--images", default="val", help="folder of images")

    # 标注文件路径
    parser.add_argument("-anno", "--annotations", default="data/val.txt", help="file of annotations")

    return parser.parse_args()

def createLabel():
    labels =[]
    for i in range(0,64,1):
        if i%2==0:
            labels.append("n_"+str(i))
        else:
            labels.append("p_"+str(i))
    labels.append("n_"+str(64))
    print(labels)
    skeleton =[]
    for i in range(1,65,2):
        skeleton.append([i,i+1])
    skeleton.append([64,65])
    print(skeleton)
    return labels,skeleton

def MainProcessing(args):
    '''main process source code.'''
    annotations = {}                                                # annotations dictionary, which will dump to json format file.
    root_path = args.root_path
    phase_folder = args.phase_folder
    images_folder = os.path.join(root_path, args.images)            # 图片文件夹 "images"
    # anno_path = os.path.join(root_path, args.annotations)           # 标注文件路径 "anno.txt"
    anno_path = args.annotations
    # coco annotations info.
    annotations["info"] = {
        "description": "customer dataset format convert to COCO format",
        "url": "http://cocodataset.org",
        "version": "1.0",
        "year": 2021,
        "contributor": "lqqq",
        "date_created": "2021"
    }
    # coco annotations licenses.
    annotations["licenses"] = [{
        "url": "https://www.apache.org/licenses/LICENSE-2.0.html",
        "id": 1,
        "name": "Apache License 2.0"
    }]
    # coco annotations categories.
    annotations["categories"] = []
    for cls, clsname in enumerate(__CLASS__):
        if clsname == '__background__':
            continue
        annotations["categories"].append(
            {
                "supercategory": "battery",
                "id": cls,
                "name": clsname
            }
        )
        for catdict in annotations["categories"]:
            if args.have_points:
                labels, skeleton = createLabel()
                catdict["keypoints"] = labels
                catdict["skeleton"] = skeleton

    for phase in phase_folder:
        annotations["images"] = []
        annotations["annotations"] = []
        if os.path.isfile(anno_path) and os.path.exists(images_folder):
            print("convert datasets {} to coco format!".format(phase))
            fd = open(anno_path, "r")
            step = 0
            for id, line in enumerate(fd.readlines()):
                if line:
                    label_info = line.split()

                    image_name = label_info[0]
                    bbox = [int(x) for x in label_info[1].split(",")]       # 标注框 bbox

                    filename = os.path.join(images_folder, image_name)

                    if os.path.isfile(filename):      # 图片文件
                        img = cv2.imread(filename)
                        height, width, _ = img.shape                         # 读取图片大小
                        x1 = bbox[0]
                        y1 = bbox[1]                                            # bbox 中心点
                        bw = bbox[2]                                            # 宽 w
                        bh = bbox[3]                                            # 高 h

                        # coco annotations images.
                        file_name = str(image_name)                                       # annotation["images"] 下 “file_name” 的值

                        annotations["images"].append(
                            {
                                "license": 1,
                                "file_name": file_name,
                                "coco_url": "data/images_crop/",
                                "height": height,
                                "width": width,
                                "date_captured": datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                                "flickr_url": "",
                                "id": id
                            }
                        )
                        # coco annotations annotations.
                        annotations["annotations"].append(
                            {
                                "id": id,
                                "image_id": id,
                                "category_id": 1,           # 类别编号，对于人体关键点检测任务，他只有一个类别，所以恒为 1
                                "segmentation": [[]],
                                "area": bw*bh,                  # 这个 area 是图像分割的东西，置为 0 就行
                                "bbox": [x1, y1, bw, bh],    # [xmin,ymin,bw,bh]
                                "iscrowd": 0,               # 目标是否被遮盖，默认为0
                            }
                        )
                        if args.have_points:
                            v = 2                           # v 字段表示关键点属性，0表示未标注，1表示已标注但不可见，2表示已标注且可见
                            catdict = annotations["annotations"][id]
                            if 1:
                                points = label_info[2].split(",")                 # [:-1]是为了去掉尾部空格
                                pp =[]
                                for i in range(0, len(points), 2):
                                    pp.append(int(points[i]))             # 字符串转换为int数值
                                    pp.append(int(points[i + 1]))
                                    pp.append(2)
                                catdict["keypoints"] = pp

                                catdict["num_keypoints"] = int(len(points) / 2)      # 一般是 65

                    step += 1
                    if step % 100 == 0:
                        print("processing {} ...".format(step))
            fd.close()
        else:
            print("WARNNING: file path incomplete, please check!")
        json_path = os.path.join(root_path+'/annotations', phase+".json")
        with open(json_path, "w") as f:
            json.dump(annotations, f)


if __name__ == "__main__":
    print("beginning to convert customer format to coco format!")
    args = argparser()
    try:
        MainProcessing(args)
    except Exception as e:
        traceback.print_exc()
    print("successful to convert customer format to coco format")
